import glob
import os
import random

import cv2
import numpy as np
import pydicom
import scipy.io

import PIL
from sklearn.cluster import k_means
from PIL import Image
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset

"""
2022年9月11日 嵇老板自己重新做了数据并成功训练了自己的方法。
"""
image_get_minmax = 0.0, 1.0

proj_get_minmax = 0.0, 4.0


def normalize(data, minmax):
    data_min, data_max = minmax
    data = np.clip(data, data_min, data_max)
    data = (data - data_min) / (data_max - data_min)
    data = data.astype(np.float32)
    data = data * 255.0
    data = np.transpose(np.expand_dims(data, 2), (2, 0, 1))
    return data


class InDuDoNetPlusDataset(Dataset):
    """
    此类专门用于InDuDoNetPlus，读取的是使用ADN中带的制作数据的代码制作的deeplesion数据
    和InDuDoNet有点不一样。
    """

    def __init__(self, config):
        len_ratio = config['len_ratio']
        self.config = config
        self.Xma_path_list = []  # 有金属尾影的图像域图
        self.Xgt_path_list = []  # 图像域的ground truth
        self.XLI_path_list = []  # 经过LI简单处理后的图像域图
        self.Sma_path_list = []  # 有金属伪影的投影域图
        self.Sgt_path_list = []  # 投影域的ground truth
        self.SLI_path_list = []  # 经过LI简单处理后的投影域图
        self.Metal_trace_list = []  # 金属轨迹正弦图
        self.data_length = 0

        # 新的成员变量 InDuDoNetPlus
        self.sigma = 1
        self.smFilter = scipy.io.loadmat(r'D:\gs\code\InDuDoNet\deeplesion\train\gaussianfilter.mat')['smFilter']
        self.miuAir = 0
        self.miuWater = 0.192
        self.starpoint = np.zeros([3, 1])
        self.starpoint[0] = self.miuAir
        self.starpoint[1] = self.miuWater
        self.starpoint[2] = 2 * self.miuWater

        self.train_mask = np.load(r"D:\gs\code\InDuDoNet\deeplesion\train\trainmask.npy")

        all_path_list = [data_dir for data_dir in glob.glob(os.path.join(config['dataroot'], "*")) if
                         os.path.isdir(data_dir)]

        for part_path in all_path_list:
            part_path = os.path.join(part_path, "train")
            for sample_path in glob.glob(os.path.join(part_path, "*")):
                for patient_path in glob.glob(os.path.join(sample_path, "*")):
                    # 有金属尾影的图像域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "ma_CT_*")))
                    self.Xma_path_list += tmp_list
                    # 图像域的ground truth
                    self.Xgt_path_list += sorted(
                        [os.path.join(patient_path, "groundtruth.mat") for i in range(len(tmp_list))])
                    # 经过LI简单处理后的图像域图
                    self.XLI_path_list += sorted(glob.glob(os.path.join(patient_path, "LI_CT_*")))
                    # 有金属伪影的投影域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "ma_sinogram_*")))
                    self.Sma_path_list += tmp_list
                    # 投影域的ground truth
                    tmp_list = sorted(
                        [os.path.join(patient_path, "groundtruth_sinogram.mat") for i in range(len(tmp_list))])
                    self.Sgt_path_list += tmp_list
                    # Sgt_path_list += tmp_list  # 加两次是因为有两种伪影
                    # 经过LI简单处理后的投影域图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "LI_sinogram_*")))
                    self.SLI_path_list += tmp_list
                    # 金属轨迹正弦图
                    tmp_list = sorted(glob.glob(os.path.join(patient_path, "metal_trace_*")))
                    self.Metal_trace_list += tmp_list

        self.Xma_path_list = self.Xma_path_list[0:int(len_ratio * len(self.Xma_path_list))]
        self.Xgt_path_list = self.Xgt_path_list[0:int(len_ratio * len(self.Xgt_path_list))]
        self.XLI_path_list = self.XLI_path_list[0:int(len_ratio * len(self.XLI_path_list))]
        self.Sma_path_list = self.Sma_path_list[0:int(len_ratio * len(self.Sma_path_list))]
        self.Sgt_path_list = self.Sgt_path_list[0:int(len_ratio * len(self.Sgt_path_list))]
        self.SLI_path_list = self.SLI_path_list[0:int(len_ratio * len(self.SLI_path_list))]
        self.Metal_trace_list = self.Metal_trace_list[0:int(len_ratio * len(self.Metal_trace_list))]

        assert len(self.Xma_path_list) == len(self.Xgt_path_list) == len(self.XLI_path_list) == len(
            self.Sma_path_list) == len(self.Sgt_path_list) == len(self.SLI_path_list) == len(self.Metal_trace_list)
        self.data_length = len(self.Xma_path_list)

    def __len__(self):
        return self.data_length

    def nmarprior(self, im, threshWater, threshBone, miuAir, miuWater, smFilter):
        imSm = scipy.ndimage.filters.convolve(im, smFilter, mode='nearest')
        # print("imSm, h:, w:", imSm.shape[0], imSm.shape[1]) # imSm, h:, w: 416 416
        priorimgHU = imSm
        priorimgHU[imSm <= threshWater] = miuAir
        h, w = imSm.shape[0], imSm.shape[1]
        priorimgHUvector = np.reshape(priorimgHU, h * w)
        region1_1d = np.where(priorimgHUvector > threshWater)
        region2_1d = np.where(priorimgHUvector < threshBone)
        region_1d = np.intersect1d(region1_1d, region2_1d)
        priorimgHUvector[region_1d] = miuWater
        priorimgHU = np.reshape(priorimgHUvector, (h, w))
        return priorimgHU

    def nmar_prior(self, XLI, M):
        XLI[M == 1] = 0.192
        h, w = XLI.shape[0], XLI.shape[1]
        im1d = XLI.reshape(h * w, 1)
        best_centers, labels, best_inertia = k_means(im1d, n_clusters=3, init=self.starpoint, max_iter=300)
        threshBone2 = np.min(im1d[labels == 2])
        threshBone2 = np.max([threshBone2, 1.2 * self.miuWater])
        threshWater2 = np.min(im1d[labels == 1])
        imPriorNMAR = self.nmarprior(XLI, threshWater2, threshBone2, self.miuAir, self.miuWater, self.smFilter)
        return imPriorNMAR

    def __getitem__(self, index):
        random_mask = random.randint(0, 89)

        Xma_path = self.Xma_path_list[index]
        Xgt_path = self.Xgt_path_list[index]
        XLI_path = self.XLI_path_list[index]
        Sma_path = self.Sma_path_list[index]
        Sgt_path = self.Sgt_path_list[index]
        SLI_path = self.SLI_path_list[index]
        Metal_trace_path = self.Metal_trace_list[index]

        # ToDo 下面变量值域的原因还没有确定，是否需要归一化也没有确定
        Xma = scipy.io.loadmat(Xma_path)['image']  #
        Xgt = scipy.io.loadmat(Xgt_path)['gt_CT']  #
        XLI = scipy.io.loadmat(XLI_path)['image']  #
        Sma = scipy.io.loadmat(Sma_path)['image']  #
        Sgt = scipy.io.loadmat(Sgt_path)['gt_sinogram'].T  # 转置是因为在制作数据时没有旋转，还不确定哪种方式更好。
        SLI = scipy.io.loadmat(SLI_path)['image']  #
        Tr = scipy.io.loadmat(Metal_trace_path)['image']

        M512 = self.train_mask[:, :, random_mask]
        M = np.array(Image.fromarray(M512).resize(Xma.shape, PIL.Image.BILINEAR))
        # ToDo 接下来的部分将和DuDoNet有些差别
        Xprior = self.nmar_prior(XLI, M)
        Xprior = normalize(Xprior, image_get_minmax)
        # cv2.imshow("M", M)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        Mask = M.astype(np.float32)
        Mask = np.transpose(np.expand_dims(Mask, 2), (2, 0, 1))

        Xma = normalize(Xma, image_get_minmax)  # Normalization function will expand the numpy array
        Xgt = normalize(Xgt, image_get_minmax)
        XLI = normalize(XLI, image_get_minmax)
        Sma = normalize(Sma, proj_get_minmax)
        Sgt = normalize(Sgt, proj_get_minmax)
        SLI = normalize(SLI, proj_get_minmax)

        Mask = np.expand_dims(Mask, 0)
        Tr = np.expand_dims(Tr, 0)

        # The order is Xma, XLI, Xgt, mask, Sma, SLI, Sgt, Tr
        return torch.Tensor(Xma), torch.Tensor(XLI), torch.Tensor(Xgt), torch.Tensor(Mask), \
               torch.Tensor(Sma), torch.Tensor(SLI), torch.Tensor(Sgt), torch.Tensor(Tr), torch.Tensor(Xprior)
